# coding=utf-8
# -*- coding: UTF-8 -*-
#
# This file is part of the kernelstudio package.
#
# (c) 2014-2025 zlin <admin@kernelstudio.com>
#
# For the full copyright and license information, please view the LICENSE file
# that was distributed with this source code.
import io
import logging
import typing as t

import pandas as pd

from ksai.document.context import DocumentSection
from ksai.document.converter.attachment import convert_attachment
from ksai.document.extractor.extractor import AbstractDocumentExtractor, DocumentExtractorException
from ksai.persistence.storage.attachment import Attachment

logger = logging.getLogger(__name__)


class ExcelDocumentExtractor(AbstractDocumentExtractor):

    def __init__(self):
        super().__init__()

    def do_extract(self, attachment: Attachment) -> t.List[DocumentSection]:
        try:
            excel_file = pd.ExcelFile(io.BytesIO(convert_attachment(attachment, 'xlsx')))
            documents = []
            for sheet_name in excel_file.sheet_names:
                try:
                    df = excel_file.parse(sheet_name=sheet_name)
                    df.dropna(how="all", inplace=True)
                    documents.append(DocumentSection(
                        title=sheet_name,
                        text=df.to_markdown(index=False),
                    ))
                except DocumentExtractorException:
                    continue
            return documents
        except Exception as e:
            raise DocumentExtractorException(f"Extract excel file error: {e}") from e

    def extensions(self) -> t.List[str]:
        return ["xls", "xlsx"]

    def mimes(self) -> t.List[str]:
        return ["application/vnd.ms-excel", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"]
